#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
工具函数模块
包含各个管理器使用的通用函数
"""

import os
import re
import subprocess
import logging
import platform
import shutil
from typing import Optional, Tuple, List, Dict, Any, Union

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("pgadmin.log")
    ]
)

logger = logging.getLogger("pgadmin")

def is_windows() -> bool:
    """检查是否在Windows系统上运行"""
    return platform.system() == "Windows"

def is_admin() -> bool:
    """检查当前用户是否有管理员权限"""
    if is_windows():
        try:
            import ctypes
            return ctypes.windll.shell32.IsUserAnAdmin() != 0
        except:
            return False
    else:
        try:
            return os.geteuid() == 0
        except AttributeError:
            # 某些Unix系统可能没有geteuid
            return False

def command_exists(command: str) -> bool:
    """检查命令是否存在"""
    return shutil.which(command) is not None

def run_command(cmd: List[str], check: bool = True, capture_output: bool = False, 
               timeout: Optional[int] = None, shell: bool = False, 
               as_user: Optional[str] = None, env: Optional[Dict[str, str]] = None) -> subprocess.CompletedProcess:
    """
    执行命令行命令并返回结果
    
    Args:
        cmd: 要执行的命令及参数
        check: 是否在命令失败时抛出异常
        capture_output: 是否捕获输出
        timeout: 命令超时时间(秒)
        shell: 是否使用shell执行
        as_user: 以指定用户身份运行命令
        env: 命令运行时的环境变量
        
    Returns:
        subprocess.CompletedProcess对象
    """
    logger.debug(f"执行命令: {' '.join(cmd)}")
    
    # 如果命令是字符串并且包含source，在Windows上需要特殊处理
    if shell and isinstance(cmd, list) and len(cmd) > 0 and cmd[0] == "source":
        if is_windows():
            # Windows上没有source命令，而是使用call
            cmd[0] = "call"
        else:
            # Linux上source是shell内建命令，需要使用.代替
            cmd[0] = "."
    
    # 如果是命令列表，检查第一个命令是否存在
    if not shell and isinstance(cmd, list) and len(cmd) > 0:
        command = cmd[0]
        if command != "sudo" and not command_exists(command):
            logger.warning(f"命令 '{command}' 不存在，执行可能会失败")
    
    # 如果指定了用户并且不是Windows系统
    if as_user and not is_windows():
        cmd = ["sudo", "-u", as_user] + cmd
    
    try:
        result = subprocess.run(
            cmd,
            check=check,
            capture_output=capture_output,
            text=capture_output,  # 如果捕获输出，则将结果作为文本返回
            timeout=timeout,
            shell=shell,
            env=env  # 添加环境变量支持
        )
        return result
    except subprocess.CalledProcessError as e:
        logger.error(f"命令执行失败: {e}")
        if check:
            raise
        # 保持返回类型一致，返回CompletedProcess而不是整数
        return subprocess.CompletedProcess(e.cmd, e.returncode, stderr=e.stderr, stdout=e.stdout)
    except subprocess.TimeoutExpired:
        logger.error(f"命令执行超时: {' '.join(cmd) if isinstance(cmd, list) else cmd}")
        raise
    except Exception as e:
        logger.error(f"执行命令时发生未知错误: {e}")
        raise

def detect_postgresql_version() -> Optional[str]:
    """
    检测并返回当前安装的PostgreSQL主要版本号
    
    Returns:
        PostgreSQL主版本号(如'16')，未检测到则返回None
    """
    # 尝试多种方法检测PostgreSQL版本
    
    # 方法1: 使用psql --version
    try:
        version_output = run_command(["psql", "--version"], capture_output=True)
        match = re.search(r"(\d+)\.\d+", version_output.stdout)
        if match:
            psql_version = match.group(1)  # 返回主版本号(如'16')
            logger.info(f"通过psql检测到PostgreSQL版本: {psql_version}")
            
            # 验证此版本服务是否存在
            if not is_windows():
                service_name = f"postgresql@{psql_version}-main.service"
                check = run_command(["systemctl", "is-enabled", service_name], check=False, capture_output=True)
                if check.returncode == 0:
                    logger.info(f"确认服务 {service_name} 存在")
                    return psql_version
                else:
                    logger.warning(f"psql版本{psql_version}的服务不存在，继续尝试其他检测方法")
            else:
                return psql_version
    except Exception as e:
        logger.warning(f"通过psql检测版本失败: {e}")
    
    # 方法2: 检查已有的postgresql服务
    if not is_windows():
        try:
            # 列出所有postgresql@*-main.service服务
            services_output = run_command(["systemctl", "list-units", "postgresql@*-main.service", "--all"], 
                                          check=False, capture_output=True)
            
            # 使用正则表达式查找服务名中的版本号
            service_matches = re.findall(r"postgresql@(\d+)-main\.service", services_output.stdout)
            if service_matches:
                service_version = service_matches[0]
                logger.info(f"通过systemd服务检测到PostgreSQL版本: {service_version}")
                return service_version
        except Exception as e:
            logger.warning(f"通过服务名检测版本失败: {e}")
    
    # 方法3: 检查常见的PostgreSQL安装路径
    common_paths = ["/usr/lib/postgresql", "/var/lib/postgresql"]
    for path in common_paths:
        if os.path.exists(path):
            try:
                # 查找版本子目录，选择数字最大的
                versions = []
                for item in os.listdir(path):
                    if item.isdigit():
                        versions.append(item)
                
                if versions:
                    # 按照数字大小排序，取最大的
                    path_version = max(versions, key=int)
                    logger.info(f"通过安装路径 {path} 检测到PostgreSQL版本: {path_version}")
                    return path_version
            except Exception as e:
                logger.warning(f"通过路径{path}检测版本失败: {e}")
    
    # 最后尝试固定的常见版本
    for version in ["16", "15", "14", "13", "12"]:
        if not is_windows():
            service_name = f"postgresql@{version}-main.service"
            try:
                check = run_command(["systemctl", "is-enabled", service_name], check=False, capture_output=True)
                if check.returncode == 0:
                    logger.info(f"找到可用的PostgreSQL服务: {service_name}")
                    return version
            except:
                pass
    
    logger.error("所有方法均无法检测到PostgreSQL版本")
    return None

def get_service_name(pg_version: Optional[str] = None) -> str:
    """
    根据PostgreSQL版本返回正确的服务名称
    
    Args:
        pg_version: PostgreSQL版本号
        
    Returns:
        服务名称
    """
    if not pg_version:
        pg_version = detect_postgresql_version()
        
    if pg_version:
        if is_windows():
            return f"postgresql-{pg_version}"
        else:
            # 检查特定版本的服务是否存在
            specific_service = f"postgresql@{pg_version}-main.service"
            try:
                # 检查服务是否存在
                systemctl_check = run_command(["systemctl", "is-enabled", specific_service], check=False, capture_output=True)
                if systemctl_check.returncode != 0:
                    # 特定版本服务不存在，返回一般的postgresql服务
                    logger.warning(f"服务 {specific_service} 不存在或未启用，尝试使用postgresql.service")
                    return "postgresql.service"
                
                # 检查服务是否活跃
                status_check = run_command(["systemctl", "is-active", specific_service], check=False, capture_output=True)
                if status_check.returncode != 0:
                    # 服务存在但不活跃，检查通用服务是否活跃
                    general_status = run_command(["systemctl", "is-active", "postgresql.service"], check=False, capture_output=True)
                    if general_status.returncode == 0:
                        logger.warning(f"服务 {specific_service} 未运行，但postgresql.service正在运行，使用通用服务")
                        return "postgresql.service"
                
                return specific_service
            except Exception as e:
                logger.error(f"检查服务状态时出错: {e}")
                return specific_service
    else:
        logger.warning("未检测到PostgreSQL版本。使用默认的'postgresql.service'服务名称")
        return "postgresql.service"

def is_dir_writable(path: str) -> bool:
    """
    检查目录是否可写
    
    Args:
        path: 目录路径
        
    Returns:
        目录是否可写
    """
    if os.path.exists(path):
        return os.access(path, os.W_OK)
    else:
        # 检查父目录是否存在且可写
        parent_dir = os.path.dirname(path)
        if not parent_dir:  # 如果是相对路径没有父目录
            parent_dir = "."
        return os.path.exists(parent_dir) and os.access(parent_dir, os.W_OK)

def confirm_action(message: str, default: bool = False) -> bool:
    """
    请求用户确认操作
    
    Args:
        message: 要显示的消息
        default: 默认选择
        
    Returns:
        用户确认结果
    """
    default_choice = "Y" if default else "N"
    choice = input(f"{message} (Y/N) [默认: {default_choice}]: ").strip().upper()
    
    if not choice:
        return default
    
    return choice in ["Y", "YES"]

def get_input_with_default(prompt: str, default: str = "") -> str:
    """
    获取用户输入，支持默认值
    
    Args:
        prompt: 提示信息
        default: 默认值
        
    Returns:
        用户输入或默认值
    """
    if default:
        result = input(f"{prompt} [默认: {default}]: ").strip()
        return result if result else default
    else:
        return input(f"{prompt}: ").strip()

def check_port_available(port: int, host: str = "127.0.0.1") -> bool:
    """
    检查指定端口是否可用
    
    Args:
        port: 要检查的端口
        host: 要检查的主机
        
    Returns:
        端口是否可用
    """
    import socket
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.settimeout(1)
        try:
            s.connect((host, port))
            return False  # 端口被占用
        except socket.error:
            return True  # 端口可用 